import unittest
from typing import override

from src.pipelines import Sequential
from src.processors import Keyword, KeywordStage, Processor, TokenizedDocumentStage


class FakeKeywordExtractor(Processor[TokenizedDocumentStage, KeywordStage]):
    @override
    def process(self, stage: TokenizedDocumentStage) -> KeywordStage:
        keywords = []
        for tokenized_document in stage.tokenized_documents:
            for word in tokenized_document:
                candidate = Keyword(word, 1.0)
                if candidate not in keywords:
                    keywords.append(candidate)
        return KeywordStage(keywords)


class TestKeywordExtractionProcessor(unittest.TestCase):
    def test_keyword_extraction_processor(self):
        stage = TokenizedDocumentStage([["hello", "world"], ["hello", "python"]])
        pipeline = Sequential(
            FakeKeywordExtractor(),
            lambda stage: stage.keywords,
        )
        self.assertEqual(
            pipeline(stage),
            [
                Keyword("hello", 1.0),
                Keyword("world", 1.0),
                Keyword("python", 1.0),
            ],
        )
